/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include "acl/acl_rt.h"
#include "acl/acl.h"
#include "aclnn_matmul_a2_b2_share_custom.h"
#include "../../../../../common/data_utils.h"
#include "../../../testcase_params.h"

namespace MatmulCustomMian {
aclrtStream CreateStream(const int32_t device)
{
    if (aclInit(nullptr) != ACL_SUCCESS) {
        printf("acl init failed.\n");
        return nullptr;
    }
    if (aclrtSetDevice(device) != ACL_SUCCESS) {
        printf("Set device failed.\n");
        CHECK_ACL(aclFinalize());
        return nullptr;
    }
    aclrtStream stream = nullptr;
    if (aclrtCreateStream(&stream) != ACL_SUCCESS) {
        printf("Create stream failed.\n");
        CHECK_ACL(aclFinalize());
        return nullptr;
    }
    return stream;
}

void DestroyStream(aclrtStream stream, const int32_t device)
{
    CHECK_ACL(aclrtDestroyStream(stream));
    if (aclrtResetDevice(device) != ACL_SUCCESS) {
        printf("Reset device failed.\n");
    }
    if (aclFinalize() != ACL_SUCCESS) {
        printf("Finalize acl failed.\n");
    }
}

struct tensorInfo {
    int64_t* dims;
    int64_t dimCnt;
    aclDataType dtype;
    aclFormat fmt;
};

void DestroyTensor(aclTensor* tensors[], void* devMem[], const int32_t tensorCount)
{
    for (auto i = 0; i < tensorCount; i++) {
        if (devMem[i]) {
            CHECK_ACL(aclrtFree(devMem[i]));
        }
        if (tensors[i]) {
            CHECK_ACL(aclDestroyTensor(tensors[i]));
        }
    }
}

uint64_t GetDataSize(struct tensorInfo* desc)
{
    if (!desc->dims) {
        return 0;
    }
    int64_t size = 1;
    for (auto i = 0; i < desc->dimCnt; i++) {
        size *= desc->dims[i];
    }
    return static_cast<uint64_t>(size * sizeof(float));
}

struct tensorIndex {
    int32_t a1 = 0;
    int32_t b1 = 1;
    int32_t a2 = 2;
    int32_t b2 = 3;
    int32_t bias = 4;
    int32_t c1 = 5;
    int32_t c2 = 6;
};

int64_t CompareResult(uint8_t* outputData, const uint64_t outSize, const std::string data_file)
{
    void* goldenData;
    CHECK_ACL(aclrtMallocHost((void**)(&goldenData), outSize));
    size_t goldenSize = static_cast<size_t>(outSize);
    bool ret = ReadFile(data_file, goldenSize, goldenData, goldenSize);
    if (ret) {
        printf("ReadFile golden success!\n");
    } else {
        CHECK_ACL(aclrtFreeHost(goldenData));
        return 1;
    }
    constexpr float EPS = 1e-4;
    int64_t wrongNum = 0;
    int64_t maxPrintNums = 10;
    for (uint64_t i = 0; i < outSize / sizeof(float); i++) {
        float a = ((float*)outputData)[i];
        float b = ((float*)goldenData)[i];
        float ae = std::abs(a - b);
        float re = ae / abs(b);
        if (ae > EPS && re > EPS) {
            printf("CompareResult index %d failed output is %lf, golden is %lf\n", i, a, b);
            wrongNum++;
        }
        if (wrongNum > maxPrintNums) {
            break;
        }
    }
    CHECK_ACL(aclrtFreeHost(goldenData));
    return wrongNum;
}

int32_t SetTensorsDevmem(aclTensor* tensors[], void* devMem[], tensorInfo tensorDesc[], const int32_t tensorCount,
    std::string dataFile[])
{
    const int32_t inputTensorCount = 5;
    for (auto i = 0; i < tensorCount; i++) {
        void* data;
        struct tensorInfo* info = &(tensorDesc[i]);
        uint64_t size = GetDataSize(info);
        if (size == 0U) {
            tensors[i] = nullptr;
            devMem[i] = nullptr;
            continue;
        }
        CHECK_ACL(aclrtMalloc(&data, size, ACL_MEM_MALLOC_HUGE_FIRST));
        // read input
        if (i < inputTensorCount) {
            size_t inputSize = size;
            void* dataHost;
            CHECK_ACL(aclrtMallocHost((void**)(&dataHost), inputSize));
            ReadFile(dataFile[i], inputSize, dataHost, inputSize);
            CHECK_ACL(aclrtMemcpy(data, size, dataHost, size, ACL_MEMCPY_HOST_TO_DEVICE));
            CHECK_ACL(aclrtFreeHost(dataHost));
        }
        devMem[i] = data;
        tensors[i] = aclCreateTensor(info->dims, info->dimCnt, info->dtype, nullptr, 0,
                                     info->fmt, info->dims, info->dimCnt, data);
    }
    return 0;
}

int32_t CheckCaseOutput(void* devMem[], tensorInfo tensorDesc[], std::string dataFile[], tensorIndex idx)
{
    int64_t wrongNum = 0;
    for (int32_t i = idx.c1; i <= idx.c2; i++) {
        uint8_t* zHost;
        uint64_t zHostSize = GetDataSize(&(tensorDesc[i]));
        CHECK_ACL(aclrtMallocHost((void**)(&zHost), zHostSize));
        CHECK_ACL(aclrtMemcpy(zHost, zHostSize, devMem[i], zHostSize, ACL_MEMCPY_DEVICE_TO_HOST));
        WriteFile(dataFile[i], zHost, zHostSize);
        wrongNum += CompareResult(zHost, zHostSize, dataFile[i]);
        CHECK_ACL(aclrtFreeHost(zHost));
    }
    if (wrongNum != 0 ) {
        printf("test failed!\n");
    } else {
        printf("test pass!\n");
    }
    return 0;
}
} // MatmulCustomMian

int32_t main(void)
{
    TestcaseParams caseParams = {24, 24 * 256, 480, 320, true};

    int64_t shapeA[] = {caseParams.m, caseParams.k};
    int64_t shapeB[] = {caseParams.k, caseParams.n};
    int64_t shapeBias[] = {caseParams.n};
    int64_t shapeC[] = {caseParams.m, caseParams.n};
    struct MatmulCustomMian::tensorInfo tensorDesc[] = {{shapeA, 2, ACL_FLOAT16, ACL_FORMAT_ND},
        {shapeB, 2, ACL_FLOAT16, ACL_FORMAT_ND}, {shapeA, 2, ACL_FLOAT16, ACL_FORMAT_ND},
        {shapeB, 2, ACL_FLOAT16, ACL_FORMAT_ND}, {shapeBias, 1, ACL_FLOAT, ACL_FORMAT_ND},
        {shapeC, 2, ACL_FLOAT, ACL_FORMAT_ND}, {shapeC, 2, ACL_FLOAT, ACL_FORMAT_ND},
    };
    aclrtStream stream = MatmulCustomMian::CreateStream(0);
    if (stream == nullptr) {
        return -1;
    }
    int32_t tensorCount = sizeof(tensorDesc) / sizeof(struct MatmulCustomMian::tensorInfo);
    aclTensor* tensors[tensorCount];
    void* devMem[tensorCount];

    constexpr MatmulCustomMian::tensorIndex idx;
    std::string dataFile[] = {"../input/a1_gm.bin", "../input/b1_gm.bin", "../input/a2_gm.bin", "../input/b2_gm.bin",
                              "../input/bias_gm.bin", "../output/output_1.bin", "../output/output_2.bin"
    };
    MatmulCustomMian::SetTensorsDevmem(tensors, devMem, tensorDesc, tensorCount, dataFile);

    size_t workspaceSize = 0;
    aclOpExecutor* handle;
    int32_t ret = 0;
    ret = aclnnMatmulA2B2ShareCustomGetWorkspaceSize(tensors[idx.a1], tensors[idx.b1], tensors[idx.a2], tensors[idx.b2],
        tensors[idx.bias], caseParams.isBias, tensors[idx.c1], tensors[idx.c2], &workspaceSize, &handle);
    if (ret != ACL_SUCCESS) {
        printf("aclnnMatmulA2B2ShareCustomGetWorkspaceSize failed. error code is %u\n", ret);
        MatmulCustomMian::DestroyTensor(tensors, devMem, tensorCount);
        MatmulCustomMian::DestroyStream(stream, 0);
        return ret;
    }
    printf("aclnnMatmulA2B2ShareCustomGetWorkspaceSize ret %u workspace size %lu\n", ret, workspaceSize);
    void* workspace = nullptr;
    if (workspaceSize != 0U) {
        CHECK_ACL(aclrtMalloc(&workspace, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST));
    }
    ret = aclnnMatmulA2B2ShareCustom(workspace, workspaceSize, handle, stream);
    printf("aclnnMatmulA2B2ShareCustom ret %u\n", ret);
    CHECK_ACL(aclrtSynchronizeStream(stream));

    MatmulCustomMian::CheckCaseOutput(devMem, tensorDesc, dataFile, idx);
    if (workspaceSize != 0U) {
        CHECK_ACL(aclrtFree(workspace));
    }

    MatmulCustomMian::DestroyTensor(tensors, devMem, tensorCount);
    MatmulCustomMian::DestroyStream(stream, 0);
    return 0;
}
